# Author Zenos
# Create 2025/2/27 上午10:08
import os
from PIL import Image
import numpy as np
import torch
from visualization import show_image_label
import json


def read_font_images(font_dir, is_train=True):
    # 读取类型先验知识
    with open(f"char_classes/{os.path.basename(font_dir)}_array.json", 'r', encoding='utf-8') as f:
        classes_dict = json.load(f)
    # 读取所有图像与标注
    txt_fname = os.path.join(font_dir, '../train.txt' if is_train else '../val.txt')
    with open(txt_fname, 'r') as f:
        images_name = f.read().split()
    images, labels, classes = [], [], []
    max_height, max_width = 0, 0  # 记录最大尺寸
    for fname in images_name:
        image = Image.open(os.path.join(font_dir, 'JPEGImagesjpg', f'{fname}.jpg')).convert("1")  # 变成二值图像
        label = Image.open(os.path.join(font_dir, 'SegmentationClassAug', f'{fname}.png'))
        image_np = np.array(image.copy()).astype(float)
        label_np = np.array(label.copy()).astype(float)
        char_classes = classes_dict[fname[2:-2]]
        # 记录最大 shape
        max_height = max(max_height, image_np.shape[0], label_np.shape[0])
        max_width = max(max_width, image_np.shape[1], label_np.shape[1])

        images.append(image_np)
        image.close()
        labels.append(np.array(label.copy()))
        label.close()
        classes.append(char_classes)
    print(f"最大尺寸: 高度={max_height}, 宽度={max_width}")
    return images, labels, classes


class FontDataset(torch.utils.data.Dataset):
    def __init__(self, is_train, font_dir):
        self.images, self.labels, self.classes = read_font_images(font_dir, is_train)

    def __getitem__(self, idx):
        p1 = 320 - self.labels[idx].shape[0]
        p2 = 320 - self.labels[idx].shape[1]
        # 标签背景填充0
        label_pad = np.pad(self.labels[idx], ((p1 // 2, p1 - p1 // 2),
                                              (p2 // 2, p2 - p2 // 2)), 'constant', constant_values=0)
        # 图像背景填充1
        feature_pad = np.pad(self.images[idx], ((
                                                    p1 // 2, p1 - p1 // 2), (p2 // 2, p2 - p2 // 2)), 'constant',
                             constant_values=1.0)
        # 类别先验信息
        classes_info = self.classes[idx]
        # classes_tensor = torch.tensor(classes_info, dtype=torch.float32)
        # # 变成形状为 [35, 1, 1] 的张量
        # classes_tensor = classes_tensor.view(35, 1, 1)
        # # 使用广播机制变成 [35, 10, 10]
        # classes_tensor = classes_tensor.expand(-1, 10, 10)
        return (torch.from_numpy(feature_pad).float().reshape([1, 320, 320]),
                torch.from_numpy(label_pad).reshape([320, 320]).long(),
                classes_info)

        # return feature_pad, label_pad

    def __len__(self):
        return len(self.images)


if __name__ == '__main__':
    TrainDataset = FontDataset(False, "/Users/zenos/Downloads/CCSSD/FZLBJW2017")
    for i in range(5):
        img, label, classes = TrainDataset[i]
        print(img.shape, label.shape, classes)
        # 转换为 numpy
        img_np = img.numpy().squeeze(0)  # feature_pad 是 (1, 320, 320)，转换后 (320, 320)
        label_np = label.numpy()  # label_pad 是 (320, 320)

        show_image_label(img_np, label_np, False)
